from typing import Union, Dict

import numpy as np


class TimeReverse:
    """
    Flip the time axis of each input
    """
    def __init__(self, keys=None) -> None:
        self.keys = keys
    
    def _apply(self, sample: np.ndarray) -> np.ndarray:
        """
        Apply time reversal to a sample.

        :param sample: numpy array with at least one dimension.
        :return: numpy array with the time axis flipped.
        """
        return np.flip(sample, axis=-1).copy()

    def _adjust_indices(self, indices: np.ndarray, length: int) -> np.ndarray:
        """
        Adjust indices such as 'rpeaks' or 'tpeaks' after the main signal has been reversed.

        :param indices: numpy array of indices (rpeaks or tpeaks).
        :param length: the total length of the time series.
        :return: adjusted numpy array of indices.
        """
        return length - 1 - indices

    def __call__(self, data: Union[np.ndarray, Dict[str, Union[np.ndarray, float]]]) -> Union[np.ndarray, Dict[str, Union[np.ndarray, float]]]:
        if isinstance(data, dict):
            for key in self.keys:
                if key == "ecg":
                    data[key] = self._apply(data[key])
                elif key == "p_offset" or key == "t_offset":
                    data[key] = self._adjust_indices(data[key], data["ecg"].shape[-1])
        elif isinstance(data, np.ndarray):
            data = self._apply(data)
        return data

class SignFlip:
    """
    Flip the sign of each input
    """
    def __init__(self, keys=None) -> None:
        self.keys = keys
    
    def _apply(self, sample: np.ndarray) -> np.ndarray:
        """
        Apply sign flip to a sample.

        :param sample: numpy array.
        :return: numpy array with the sign of all elements flipped.
        """
        return -sample

    def __call__(self, data: Union[np.ndarray, Dict[str, Union[np.ndarray, float]]]) -> Union[np.ndarray, Dict[str, Union[np.ndarray, float]]]:
        if isinstance(data, dict):
            for key in self.keys:
                if key in data:
                    data[key] = self._apply(data[key])
        elif isinstance(data, np.ndarray):
            data = self._apply(data)
        return data